#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <time.h>
#include <signal.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <linux/perf_event.h>
#include <bpf/libbpf.h>
#include <bpf/bpf.h>
#include "cpu_monitor.skel.h"

// 与内核共享的事件结构体
struct event {
    unsigned int pid;
    unsigned int tid;
    unsigned long stack_id;
    char comm[16];
};

// 进程级统计
struct proc_stats {
    unsigned long count;
    char comm[16];
};

// 线程级统计
struct thread_stats {
    unsigned long count;
    char comm[16];
};

// 函数级统计
struct func_stats {
    unsigned long count;
    char func[64];
};

// 全局变量
static volatile bool exiting = false;
static struct proc_stats *proc_map = NULL;
static struct thread_stats *thread_map = NULL;
static struct func_stats *func_map = NULL;
static int stack_map_fd = -1;
static unsigned long total_samples = 0;

static int libbpf_print_fn(enum libbpf_print_level lvl, const char *fmt, va_list args) {
    return vfprintf(stderr, fmt, args);
}

static void sig_handler(int sig) {
    exiting = true;
}

static void print_report() {
    printf("\n=== CPU Usage Report (Total Samples: %lu) ===\n", total_samples);
    
    // 进程级报告
    printf("\nTop Processes:\n");
    printf("%-8s %-16s %-10s %s\n", "PID", "COMM", "SAMPLES", "PERCENT");
    for (int i = 0; i < 1024; i++) {
        if (proc_map[i].count > 0) {
            printf("%-8d %-16s %-10lu %.2f%%\n", 
                   i, proc_map[i].comm, 
                   proc_map[i].count,
                   (proc_map[i].count * 100.0) / total_samples);
        }
    }
    
    // 函数级报告
    printf("\nTop Functions:\n");
    printf("%-64s %-10s %s\n", "FUNCTION", "SAMPLES", "PERCENT");
    for (int i = 0; i < 32768; i++) {
        if (func_map[i].count > 10) { // 只显示有意义的样本
            printf("%-64s %-10lu %.2f%%\n", 
                   func_map[i].func, 
                   func_map[i].count,
                   (func_map[i].count * 100.0) / total_samples);
        }
    }
}

static void handle_event(void *ctx, int cpu, void *data, __u32 sz) {
    struct event *e = data;
    total_samples++;
    
    // 更新进程级统计
    if (e->pid < 1024) {
        proc_map[e->pid].count++;
        strncpy(proc_map[e->pid].comm, e->comm, sizeof(proc_map[e->pid].comm));
    }
    
    // 更新线程级统计
    if (e->tid < 4096) {
        thread_map[e->tid].count++;
        strncpy(thread_map[e->tid].comm, e->comm, sizeof(thread_map[e->tid].comm));
    }
    
    // 更新函数级统计
    if (e->stack_id != (unsigned long)-1) {
        unsigned int hash = e->stack_id % 32768;
        func_map[hash].count++;
        snprintf(func_map[hash].func, sizeof(func_map[hash].func), 
                 "StackID:%lu (PID:%d)", e->stack_id, e->pid);
    }
}

static void handle_lost_events(void *ctx, int cpu, __u64 cnt) {
    fprintf(stderr, "Lost %llu events on CPU %d\n", cnt, cpu);
}

int main(int argc, char **argv) {
    struct perf_buffer *pb = NULL;
    struct cpu_monitor_bpf *skel;
    int err, ncpus;
    
    // 设置日志回调
    libbpf_set_print(libbpf_print_fn);
    
    // 设置信号处理
    signal(SIGINT, sig_handler);
    signal(SIGTERM, sig_handler);
    
    // 分配统计内存
    proc_map = calloc(1024, sizeof(struct proc_stats));
    thread_map = calloc(4096, sizeof(struct thread_stats));
    func_map = calloc(32768, sizeof(struct func_stats));
    if (!proc_map || !thread_map || !func_map) {
        fprintf(stderr, "Failed to allocate stats memory\n");
        err = -ENOMEM;
        goto cleanup;
    }
    
    // 1. 打开并加载BPF程序
    skel = cpu_monitor_bpf__open();
    if (!skel) {
        fprintf(stderr, "Failed to open BPF skeleton\n");
        err = -1;
        goto cleanup;
    }
    
    // 2. 加载BPF程序
    if ((err = cpu_monitor_bpf__load(skel))) {
        fprintf(stderr, "Failed to load BPF skeleton: %d\n", err);
        goto cleanup;
    }
    
    // 获取堆栈映射文件描述符
    stack_map_fd = bpf_map__fd(skel->maps.stack_traces);
    
    // 3. 设置perf事件采样频率 (99Hz)
    struct perf_event_attr attr = {
        .type = PERF_TYPE_SOFTWARE,
        .size = sizeof(attr),
        .config = PERF_COUNT_SW_CPU_CLOCK,
        .sample_freq = 99,
        .freq = 1,
    };
    
    // 4. 附加到所有CPU
    ncpus = libbpf_num_possible_cpus();
    for (int cpu = 0; cpu < ncpus; cpu++) {
        int fd = syscall(__NR_perf_event_open, &attr, -1, cpu, -1, 0);
        if (fd < 0) {
            perror("perf_event_open failed");
            err = -errno;
            goto cleanup;
        }
        if (bpf_program__attach_perf_event(skel->progs.profile_cpu, fd) < 0) {
            perror("attach_perf_event failed");
            close(fd);
            err = -errno;
            goto cleanup;
        }
    }
    
    // 5. 设置perf缓冲区 - 使用正确的映射类型
    pb = perf_buffer__new(bpf_map__fd(skel->maps.events), 64, 
                         handle_event, handle_lost_events, NULL, NULL);
    if (!pb) {
        err = -errno;
        fprintf(stderr, "Failed to create perf buffer: %d\n", err);
        goto cleanup;
    }
    
    printf("Monitoring CPU usage. Press Ctrl+C to exit and show report.\n");
    
    // 6. 事件处理循环
    while (!exiting) {
        err = perf_buffer__poll(pb, 1000);
        if (err < 0 && err != -EINTR) {
            fprintf(stderr, "Error polling perf buffer: %d\n", err);
            break;
        }
        print_report();
    }
    
    // 打印最终报告
    // print_report();

cleanup:
    // 清理资源
    if (pb) perf_buffer__free(pb);
    if (skel) cpu_monitor_bpf__destroy(skel);
    if (proc_map) free(proc_map);
    if (thread_map) free(thread_map);
    if (func_map) free(func_map);
    return err < 0 ? -err : 0;
}